#!/usr/bin/env python3
# -*- coding: utf-8 -*-
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel.ext.asyncio.session import AsyncSession
from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlalchemy.pool import NullPool, AsyncAdaptedQueuePool, QueuePool
from sqlmodel import SQLModel, create_engine
from collections.abc import AsyncGenerator
import logging

from settings.config import settings
from .base import BaseModel

logger = logging.getLogger(__name__)

# DATABASE_URI = (
#     f'mysql+pymysql://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:'
#     f'{settings.DATABASE_PORT}/{settings.DATABASE_NAME}?charset={settings.DATABASE_CHARSET}'
# )

# engine = create_engine(url=DATABASE_URI, echo=settings.DATABASE_ECHO)

# connect_args = {"check_same_thread": False}

ASYNC_DATABASE_URI = (
    f'mysql+asyncmy://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:'
    f'{settings.DATABASE_PORT}/{settings.DATABASE_NAME}?charset={settings.DATABASE_CHARSET}'
)

SYNC_DATABASE_URI = (
    f'mysql+pymysql://{settings.DATABASE_USER}:{settings.DATABASE_PASSWORD}@{settings.DATABASE_HOST}:'
    f'{settings.DATABASE_PORT}/{settings.DATABASE_NAME}?charset={settings.DATABASE_CHARSET}'
)

async_engine = create_async_engine(
    url=ASYNC_DATABASE_URI,
    echo=settings.DATABASE_ECHO,
    poolclass=AsyncAdaptedQueuePool,
    pool_size=settings.DATABASE_POOL_SIZE,
    # max_overflow=64,
)

sync_engine = create_engine(
    url=SYNC_DATABASE_URI,
    echo=settings.DATABASE_ECHO,
    poolclass=QueuePool,
    pool_size=settings.DATABASE_POOL_SIZE,
)

SessionLocal = async_sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=async_engine,
    class_=AsyncSession,
    expire_on_commit=False
)   

async def get_db() -> AsyncGenerator[AsyncSession, None]:
    async with SessionLocal() as session:
        try:
            yield session
        except Exception as e:
            await session.rollback()
            raise e
        finally:
            await session.close()

async def init_db():
    """初始化数据库连接
    """
    logger.info("初始化数据库连接...")
    async with async_engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
    
    logger.info("数据库连接成功")